package org.skipper.dubbo.spring.starter.model;

import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author shoujian
 * Date: 2023-07-27
 * Time: 21:14
 */
public class RpcTraceContext implements Serializable {

    private final static ThreadLocal<Map<String, Object>> traceContextHolder = new ThreadLocal<>();


    public static ThreadLocal<Map<String, Object>> get() {
        return traceContextHolder;
    }

    /**
     * 设置traceContext
     *
     * @param traceContext traceContext
     */
    public static void setAttachments(Map<String, Object> traceContext) {
        traceContextHolder.set(traceContext);
    }

    /**
     * 获取traceContext
     *
     * @return traceContext
     */
    public static Map<String, Object> getTraceContext() {
        return traceContextHolder.get();
    }

    /**
     * 清空trace上下文
     */
    public static void clear() {
        traceContextHolder.remove();
    }

    /**
     * 新增 key-value 到 traceContext
     *
     * @param key   the key to add
     * @param value the value to add
     */
    public static void put(String key, Object value) {
        Map<String, Object> traceContext = traceContextHolder.get();
        if (traceContext == null) {
            traceContext = new HashMap<>();
            traceContextHolder.set(traceContext);
        }
        traceContext.put(key, value);
    }

}